from typing import Union, Dict, List
from pathlib import Path
import json
import os

from .utils import loads_json, load_json


def group_entity_by_type(dataset):
    """_summary_

    Args:
        dataset (_type_): [{
        "sentence": "An art exhibit at the Hakawati Theatre in Arab east Jerusalem was a series of portraits of Palestinians killed in the rebellion .",
        "relations": [
            {
                "head": {
                    "name": "Hakawati Theatre",
                    "type": "Organization",
                    "pos": [22, 38],
                },
                "type": "Organization based in",
                "tail": {"name": "Jerusalem", "type": "Location", "pos": [52, 61]},
            }
        ],
    }]

    Returns:
        {
            "Organization": ["Hakawati Theatre"],
            "Location": ["Jerusalem"],
        }
    """
    entity_type_dict = {}
    for data_item in dataset:
        for r in data_item["relations"]:
            head = r["head"]
            tail = r["tail"]
            if head["type"] not in entity_type_dict:
                entity_type_dict[head["type"]] = []
            if tail["type"] not in entity_type_dict:
                entity_type_dict[tail["type"]] = []
            entity_type_dict[head["type"]].append(head["name"])
            entity_type_dict[tail["type"]].append(tail["name"])

    return entity_type_dict


def save_entity_by_type_from_train_path(
    train_path: Union[str, Path], save_path: Union[str, Path]
) -> None:
    """从 train.json 中根据实体类型，并保存为 entity_by_type.json

    Args:
        train_path (Union[str, Path]): 训练集路径，用于提取实体类型
        save_path (Union[str, Path]): 保存路径
    """
    data = load_json(train_path)
    entity_by_type = group_entity_by_type(data)
    with open(save_path, "w") as f:
        json.dump(entity_by_type, f)


def get_entity_by_type_json(
    root_dir: Union[str, Path], dataset_name: Union[str, Path]
) -> Dict:
    """获取数据集中entity_by_type, 根据实体类型分类的所有实体

    Args:
        dataset_name (Union[str, Path]): 数据集名称

    Returns:
        Dict: 返回数据集中的实体类型
    """
    data_dir = Path(root_dir) / dataset_name
    entity_by_type_path = f"{data_dir}/entity_by_type.json"

    # 判断 entity_by_type_path 是否存在， 如果存在，读取并返回，如果不存在，从 train.json 中提取实体类型
    if not os.path.exists(entity_by_type_path):
        save_entity_by_type_from_train_path(
            f"{data_dir}/train.json", entity_by_type_path
        )
    return load_json(entity_by_type_path)


def save_relations_by_type_from_train_path(
    train_path: Union[str, Path], save_path: Union[str, Path]
) -> None:
    """从 train.json 中提取关系类型，并保存为 relations_by_type.json

    Args:
        train_path (Union[str, Path]): 训练集路径，用于提取关系类型
        save_path (Union[str, Path]): 保存路径
    """
    data = load_json(train_path)
    relations_by_type = {}
    for sample in data:
        for rel in sample["relations"]:
            rel_type = rel["type"]
            if rel_type not in relations_by_type:
                relations_by_type[rel_type] = []
            relations_by_type[rel_type].append(
                {
                    "head": {"name": rel["head"]["name"], "type": rel["head"]["type"]},
                    "tail": {"name": rel["tail"]["name"], "type": rel["tail"]["type"]},
                    "type": rel_type,
                }
            )
    with open(save_path, "w") as f:
        json.dump(relations_by_type, f)


def get_relations_by_type_json(
    root_dir: Union[str, Path], dataset_name: Union[str, Path]
) -> Dict:
    """获取数据集中relations_by_type, 根据关系类型分类的所有关系三元组

    Args:
        dataset_name (Union[str, Path]): 数据集名称

    Returns:
        Dict: 返回数据集中的关系类型
    """
    data_dir = Path(root_dir) / dataset_name
    relations_by_type_path = f"{data_dir}/relations_by_type.json"

    # 判断 entity_by_type_path 是否存在， 如果存在，读取并返回，如果不存在，从 train.json 中提取实体类型
    if not os.path.exists(relations_by_type_path):
        save_relations_by_type_from_train_path(
            f"{data_dir}/train.json", relations_by_type_path
        )
    return load_json(relations_by_type_path)
